#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
增强版 MEKF（带鲁棒处理）并加入逐步日志打印与运行诊断，方便判断滤波器计算是否合理。
用法：在命令行运行该文件（或在 IDE 中运行），会在当前工作目录生成日志文件 `mekf_run.log`。
要点：每 step 会按设定频率打印关键量（可配置），并在结束时输出诊断汇总。
"""

import numpy as np
import logging
import datetime
import os

# -------------------- logging configuration --------------------
LOG_FILE = 'mekf_run.log'
# If you want a very verbose log (every step), set LOG_EVERY = 1. For long sims choose >10.
LOG_EVERY = 20
VERBOSE = True  # console prints when True; file always logs
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s %(levelname)s: %(message)s',
    handlers=[
        logging.FileHandler(LOG_FILE, mode='w', encoding='utf-8'),
        logging.StreamHandler() if VERBOSE else logging.NullHandler()
    ]
)
logger = logging.getLogger('MEKF')

# -------------------- quaternion utilities --------------------
def skew(v):
    return np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])

def quat_mul(q1, q2):
    w1, x1, y1, z1 = q1
    w2, x2, y2, z2 = q2
    return np.array([
        w1*w2 - x1*x2 - y1*y2 - z1*z2,
        w1*x2 + x1*w2 + y1*z2 - z1*y2,
        w1*y2 - x1*z2 + y1*w2 + z1*x2,
        w1*z2 + x1*y2 - y1*x2 + z1*w2
    ])

def quat_conjugate(q):
    w, x, y, z = q
    return np.array([w, -x, -y, -z])

def normalize_quat(q):
    n = np.linalg.norm(q)
    if n <= 0:
        return np.array([1.0, 0.0, 0.0, 0.0])
    return q / n

def small_angle_quat(dtheta):
    angle = np.linalg.norm(dtheta)
    if angle < 1e-12:
        return np.array([1.0, 0.0, 0.0, 0.0])
    axis = dtheta / angle
    w = np.cos(angle/2.0)
    xyz = axis * np.sin(angle/2.0)
    return np.concatenate([[w], xyz])

def integrate_quat(q, omega, dt):
    dtheta = omega * dt
    dq = small_angle_quat(dtheta)
    q_new = quat_mul(dq, q)
    return normalize_quat(q_new)

def quat_to_euler(q):
    w, x, y, z = q
    sinr_cosp = 2.0*(w*x + y*z)
    cosr_cosp = 1.0 - 2.0*(x*x + y*y)
    roll = np.arctan2(sinr_cosp, cosr_cosp)
    sinp = 2.0*(w*y - z*x)
    if abs(sinp) >= 1:
        pitch = np.sign(sinp) * (np.pi/2)
    else:
        pitch = np.arcsin(sinp)
    siny_cosp = 2.0*(w*z + x*y)
    cosy_cosp = 1.0 - 2.0*(y*y + z*z)
    yaw = np.arctan2(siny_cosp, cosy_cosp)
    return np.array([roll, pitch, yaw])

def quat_to_rotvec(q):
    qn = q / np.linalg.norm(q)
    w = np.clip(qn[0], -1.0, 1.0)
    angle = 2.0 * np.arccos(w)
    if angle < 1e-12:
        return np.zeros(3)
    s = np.sqrt(1 - w*w)
    axis = qn[1:]/(s + 1e-12)
    return axis * angle

# -------------------- simulation params (same as before) --------------------
dt = 0.01
T = 20.0
N = int(T / dt)
t = np.arange(N) * dt

omega_true = np.zeros((N, 3))
omega_true[:, 2] = 0.3
omega_true[:, 0] = 0.12 * np.sin(0.6 * t)
omega_true[:, 1] = 0.10 * np.cos(0.4 * t)

g_inertial = np.array([0.0, 0.0, 1.0])
m_inertial = np.array([0.15, 0.0, 0.9887])
m_inertial = m_inertial / np.linalg.norm(m_inertial)

gyro_bias_true = np.array([0.02, -0.015, 0.01])
gyro_noise_std = 0.004
gyro_bias_walk_std = 5e-5
accel_noise_std = 0.02
mag_noise_std = 0.01

# generate ground truth and noisy sensors like previous script
qs_true = np.zeros((N,4))
q_true = np.array([1.,0.,0.,0.])
qs_true[0]=q_true
for k in range(1,N):
    q_true = integrate_quat(q_true, omega_true[k-1], dt)
    qs_true[k]=q_true

gyro_meas = np.zeros((N,3))
accel_meas = np.zeros((N,3))
mag_meas = np.zeros((N,3))

for k in range(N):
    qk = qs_true[k]
    Rb_i = quat_to_euler(qk)  # not used for measurements beyond true quat
    true_accel_b = quat_to_euler(qk)  # placeholder (we will compute properly below)

# regenerate measurements with linear accel and magnetic disturbances (clean, same logic)
qs_true = np.zeros((N,4))
q_true = np.array([1.,0.,0.,0.])
qs_true[0]=q_true
for k in range(1,N):
    q_true = integrate_quat(q_true, omega_true[k-1], dt)
    qs_true[k]=q_true

for k in range(N):
    qk = qs_true[k]
    Rb_i = np.array([  # compute rotation matrix via small helper to avoid dependency
        [1-2*(qk[2]**2+qk[3]**2),   2*(qk[1]*qk[2] - qk[3]*qk[0]),  2*(qk[1]*qk[3] + qk[2]*qk[0])],
        [2*(qk[1]*qk[2] + qk[3]*qk[0]),   1-2*(qk[1]**2+qk[3]**2),  2*(qk[2]*qk[3] - qk[1]*qk[0])],
        [2*(qk[1]*qk[3] - qk[2]*qk[0]),   2*(qk[2]*qk[3] + qk[1]*qk[0]),  1-2*(qk[1]**2+qk[2]**2)]
    ])
    true_accel_b = Rb_i.T @ g_inertial
    true_mag_b = Rb_i.T @ m_inertial
    gyro_meas[k] = omega_true[k] + gyro_bias_true + np.random.randn(3)*gyro_noise_std

    # linear acceleration events
    if 5.0 < t[k] < 6.0 or 14.5 < t[k] < 15.5:
        lin_acc = np.array([0.6*np.sin(10*t[k]), 0.4*np.cos(8*t[k]), 0.0])
    else:
        lin_acc = np.zeros(3)
    accel_vec = true_accel_b + lin_acc + np.random.randn(3)*accel_noise_std
    accel_meas[k] = accel_vec / (np.linalg.norm(accel_vec)+1e-12)

    if 9.8 < t[k] < 10.4 or 15.2 < t[k] < 15.4:
        mag_disturb = np.array([0.8, 0.8, -0.6])
    else:
        mag_disturb = np.zeros(3)
    mag_vec = true_mag_b + mag_disturb + np.random.randn(3)*mag_noise_std
    mag_meas[k] = mag_vec / (np.linalg.norm(mag_vec)+1e-12)

# -------------------- MEKF with logging --------------------
q_est = np.array([1.,0.,0.,0.])
b_est = np.zeros(3)

P = np.zeros((6,6))
P[0:3,0:3] = np.eye(3) * (np.deg2rad(3.0)**2)
P[3:6,3:6] = np.eye(3) * (0.03**2)

gyro_noise_var = gyro_noise_std**2
bias_rw_var = gyro_bias_walk_std**2
Q = np.zeros((6,6))
Q[0:3,0:3] = np.eye(3) * gyro_noise_var * dt
Q[3:6,3:6] = np.eye(3) * bias_rw_var * dt

R_acc_base = np.eye(3) * (accel_noise_std**2)
R_mag_base = np.eye(3) * (mag_noise_std**2)

accel_norm_window = 0.5
window_len = max(1, int(accel_norm_window/dt))
accel_norm_ma = 1.0
ma_alpha = 1.0/window_len
mag_angle_threshold = np.deg2rad(30.0)
accel_norm_threshold = 0.25
mahal_threshold = 20.0

q_est_store = np.zeros((N,4))
b_est_store = np.zeros((N,3))
rotvec_errors = np.zeros((N,3))
accel_skip = np.zeros(N, dtype=bool)
mag_skip = np.zeros(N, dtype=bool)

logger.info("MEKF run start: dt=%.4f, N=%d", dt, N)
start_time = datetime.datetime.now()

for k in range(N):
    omega_m = gyro_meas[k]
    omega_unbiased = omega_m - b_est
    q_est = integrate_quat(q_est, omega_unbiased, dt)

    # state transition and covariance predict
    F = np.zeros((6,6))
    F[0:3,0:3] = -skew(omega_unbiased)
    F[0:3,3:6] = -np.eye(3)
    Phi = np.eye(6) + F*dt
    P = Phi @ P @ Phi.T + Q

    z_acc = accel_meas[k]
    z_mag = mag_meas[k]

    Rb_i_est = np.array([
        [1-2*(q_est[2]**2+q_est[3]**2),   2*(q_est[1]*q_est[2] - q_est[3]*q_est[0]),  2*(q_est[1]*q_est[3] + q_est[2]*q_est[0])],
        [2*(q_est[1]*q_est[2] + q_est[3]*q_est[0]),   1-2*(q_est[1]**2+q_est[3]**2),  2*(q_est[2]*q_est[3] - q_est[1]*q_est[0])],
        [2*(q_est[1]*q_est[3] - q_est[2]*q_est[0]),   2*(q_est[2]*q_est[3] + q_est[1]*q_est[0]),  1-2*(q_est[1]**2+q_est[2]**2)]
    ])
    h_acc = Rb_i_est.T @ g_inertial
    h_mag = Rb_i_est.T @ m_inertial

    accel_norm = np.linalg.norm(accel_meas[k])
    accel_norm_ma = (1-ma_alpha)*accel_norm_ma + ma_alpha*accel_norm
    accel_deviation = abs(accel_norm - 1.0)

    if accel_deviation > accel_norm_threshold:
        accel_weight = 100.0
        accel_skip[k] = True
    elif accel_deviation > (accel_norm_threshold*0.6):
        accel_weight = 10.0
    else:
        accel_weight = 1.0

    cosang = np.clip(np.dot(z_mag, h_mag)/(np.linalg.norm(z_mag)*np.linalg.norm(h_mag)+1e-12), -1.0, 1.0)
    mag_angle = np.arccos(cosang)
    if mag_angle > mag_angle_threshold:
        mag_weight = 100.0
        mag_skip[k] = True
    else:
        mag_weight = 1.0

    H_acc = np.hstack([ -skew(h_acc), np.zeros((3,3)) ])
    H_mag = np.hstack([ -skew(h_mag), np.zeros((3,3)) ])
    H = np.vstack([H_acc, H_mag])
    z = np.hstack([z_acc, z_mag])
    h = np.hstack([h_acc, h_mag])

    R_block = np.zeros((6,6))
    R_block[0:3,0:3] = R_acc_base * accel_weight
    R_block[3:6,3:6] = R_mag_base * mag_weight

    y = z - h
    S = H @ P @ H.T + R_block
    try:
        S_inv = np.linalg.inv(S)
    except np.linalg.LinAlgError:
        S_inv = np.linalg.pinv(S)
    maha = float(y.T @ S_inv @ y)

    # Adaptive gating (same logic)
    if maha > mahal_threshold:
        if mag_angle > mag_angle_threshold:
            R_block[3:6,3:6] *= 1e3
            S = H @ P @ H.T + R_block
            try:
                S_inv = np.linalg.inv(S)
            except np.linalg.LinAlgError:
                S_inv = np.linalg.pinv(S)
            maha = float(y.T @ S_inv @ y)
    if maha > mahal_threshold:
        R_block[0:3,0:3] *= 1e2
        S = H @ P @ H.T + R_block
        try:
            S_inv = np.linalg.inv(S)
        except np.linalg.LinAlgError:
            S_inv = np.linalg.pinv(S)
        maha = float(y.T @ S_inv @ y)

    skipped = False
    if maha > (mahal_threshold*2.0):
        skipped = True

    # Logging key quantities at configured frequency
    if (k % LOG_EVERY == 0) or skipped:
        # small summaries to help judge filter behavior
        qnorm = np.linalg.norm(q_est)
        P_trace = np.trace(P)
        P_diag = np.diag(P)
        K_approx = None
        try:
            K_approx = P @ H.T @ np.linalg.inv(S + 1e-12*np.eye(S.shape[0]))
            K_norm = np.linalg.norm(K_approx)
        except Exception:
            K_norm = np.nan

        logger.info("STEP %d time=%.3f s | qnorm=%.6f | maha=%.3f | mag_ang=%.2f deg | accel_dev=%.3f | accel_w=%.1f mag_w=%.1f | skipped=%s",
                    k, t[k], qnorm, maha, np.degrees(mag_angle), accel_deviation, accel_weight, mag_weight, skipped)
        logger.debug("    trace(P)=%.6e P_diag=[%.3e, %.3e, %.3e, %.3e, %.3e, %.3e] K_norm=%.6e b_est=[%.4e, %.4e, %.4e]",
                     P_trace, P_diag[0], P_diag[1], P_diag[2], P_diag[3], P_diag[4], P_diag[5], K_norm, b_est[0], b_est[1], b_est[2])
        logger.debug("    innovation_norm=%.6e | innovation(ang acc part)=[%.3e,%.3e,%.3e] | innovation(mag part)=[%.3e,%.3e,%.3e]",
                     np.linalg.norm(y), y[0], y[1], y[2], y[3], y[4], y[5])

    if skipped:
        # skip measurement update
        q_est_store[k] = q_est
        b_est_store[k] = b_est
        q_est_conj = quat_conjugate(q_est)
        q_rel = quat_mul(q_est_conj, qs_true[k])
        rotvec_errors[k] = quat_to_rotvec(q_rel)
        continue

    # Kalman gain and update
    K = P @ H.T @ np.linalg.inv(S + 1e-12*np.eye(S.shape[0]))
    dx = K @ y
    dtheta = dx[0:3]
    db = dx[3:6]
    dq = small_angle_quat(dtheta)
    q_est = normalize_quat(quat_mul(dq, q_est))
    b_est = b_est + db

    I6 = np.eye(6)
    P = (I6 - K @ H) @ P @ (I6 - K @ H).T + K @ R_block @ K.T

    q_est_store[k] = q_est
    b_est_store[k] = b_est
    q_est_conj = quat_conjugate(q_est)
    q_rel = quat_mul(q_est_conj, qs_true[k])
    rotvec_errors[k] = quat_to_rotvec(q_rel)

end_time = datetime.datetime.now()
elapsed = (end_time - start_time).total_seconds()

# final diagnostics and summary logging
rms_rotvec = np.sqrt(np.mean(rotvec_errors**2, axis=0))
rms_rotvec_norm = np.sqrt(np.mean(np.sum(rotvec_errors**2,axis=1)))
logger.info("MEKF run finished in %.3f s (wall). RMS rotvec components (rad) x=%.4f y=%.4f z=%.4f",
            elapsed, rms_rotvec[0], rms_rotvec[1], rms_rotvec[2])
logger.info("RMS rotvec norm (rad)=%.4f deg=%.3f", rms_rotvec_norm, rms_rotvec_norm*180/np.pi)
logger.info("Fraction accel skipped/downweighted=%.3f mag skipped/downweighted=%.3f",
            np.mean(accel_skip.astype(float)), np.mean(mag_skip.astype(float)))
logger.info("Final P trace=%.6e | P diag (final) = %s", np.trace(P), np.array2string(np.diag(P), precision=3, separator=','))
try:
    eigs = np.linalg.eigvals(P)
    logger.info("Final P eigenvalues: %s", np.array2string(eigs, precision=3, separator=','))
except Exception as e:
    logger.warning("Failed to compute eigenvalues of P: %s", str(e))

logger.info("Log file: %s  (current working directory: %s)", os.path.abspath(LOG_FILE), os.getcwd())
print("Done. See log file 'mekf_run.log' for step-by-step details.")
